#include <stdio.h>
#include <stdint.h>
#include <assert.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include "comdef.h"

/* 前16个字节为数据头 */
#define IMAGE_HEAD_LEN (16l)
#define LABEL_HEAD_LEN (8l)
#define TRAIN_NUM (60000)
#define TEST_NUM  (10000)
#define CLASS_NUM (10)
#define BATCH_SIZE (10)
#define IMAGE_ROW (28)
#define IMAGE_COL (28)
#define NUM_INPUT (IMAGE_ROW*IMAGE_COL)
#define NUM_OUTPUT (10)
#define ETA (0.003)
#define NUM_EPOCHS (160000)
#define NUM_EPOCHS_TEST (1000000)

typedef enum class_name_e
{
    t_shirt, 
    trouser, 
    pullover, 
    dress, 
    coat,
    sandal, 
    shirt, 
    sneaker, 
    bag, 
    ankle_boot,
}class_name_m;

const char train_images_path[] = "data/train-images-idx3-ubyte";
const char train_labels_path[] = "data/train-labels-idx1-ubyte";
const char test_images_path[] = "data/t10k-images-idx3-ubyte";
const char test_labels_path[] = "data/t10k-labels-idx1-ubyte";
const char class_name[10][20] = 
{
    "t-shirt", 
    "trouser", 
    "pullover", 
    "dress", 
    "coat",
    "sandal", 
    "shirt", 
    "sneaker", 
    "bag", 
    "ankle boot",
};

void dump_file(FILE *fp, uint64_t size)
{
    uint8_t buf[1];
    uint64_t len = 0;
    assert(fseek(fp, 0, SEEK_SET) == 0);
    uint64_t read_len = fread(buf, sizeof(buf), 1, fp);

    while(read_len > 0)
    {
        if(len >= size)
        {
            printf("\n");
            return;
        }
        printf("%02x ", buf[0]);
        if((len + 1) % 16 == 0)
        {
            printf("\n");
        }
        len++;
        read_len = fread(buf, sizeof(buf), 1, fp);
    }
    printf("\n");
    return ;
}


uint8_t *get_one_image(
    uint8_t *image_mem, 
    uint64_t image_index)
{
    uint8_t *image_addr = (uint8_t *)((uint64_t)image_mem + IMAGE_HEAD_LEN + image_index*NUM_INPUT);
    return image_addr;
}

uint64_t get_one_label(
    uint8_t *label_mem, 
    uint64_t image_index)
{
    uint8_t *label_addr = (uint8_t *)((uint64_t)label_mem + LABEL_HEAD_LEN + image_index);
    return (uint64_t)*label_addr;
}


void softmax(
    double *input_data, 
    uint64_t row, 
    uint64_t col, 
    double *output_data)
{
    for(uint64_t index = 0;index < (row*col);index++)
    {
        output_data[index] = exp(input_data[index]);
    }
    double *partition = (double *)malloc(sizeof(double)*row);
    for(uint64_t i = 0;i<row;i++)
    {
        partition[i] = 0;
        for(uint64_t j=0;j<col;j++)
        {
            partition[i] += output_data[i*col+j];
        }
    }

    for(uint64_t i = 0;i<row;i++)
    {
        for(uint64_t j=0;j<col;j++)
        {
            output_data[i*col+j] = output_data[i*col+j] / partition[i];
        }
    }
}

void print_actual_w(double *actual_w)
{
    for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
    {
        for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
        {
            printf("actual_w[%ld][%ld]:%f\n", output_index, input_index, actual_w[output_index * NUM_INPUT + input_index]);
        }
    }
}

int softmax_new(void)
{
    FUNC_START;
    uint64_t flen = 0;
    FILE *train_images_fp = NULL;
    train_images_fp = fopen(train_images_path, "r");
    fseek(train_images_fp,0L,SEEK_END);
    flen=ftell(train_images_fp);
    uint8_t *train_images_mem = (uint8_t *)malloc(flen+1);
    assert(train_images_mem != NULL);
    fseek(train_images_fp,0L,SEEK_SET);
    fread(train_images_mem, flen, 1, train_images_fp);
    fclose(train_images_fp);

    FILE *train_labels_fp = NULL;
    train_labels_fp = fopen(train_labels_path, "r");
    fseek(train_labels_fp,0L,SEEK_END);
    flen=ftell(train_labels_fp);
    uint8_t *train_labels_mem = (uint8_t *)malloc(flen+1);
    assert(train_labels_mem != NULL);
    fseek(train_labels_fp,0L,SEEK_SET);
    fread(train_labels_mem, flen, 1, train_images_fp);
    fclose(train_labels_fp);

    FILE *test_images_fp = NULL;
    test_images_fp = fopen(test_images_path, "r");
    fseek(test_images_fp,0L,SEEK_END);
    flen=ftell(test_images_fp);
    uint8_t *test_images_mem = (uint8_t *)malloc(flen+1);
    assert(test_images_mem != NULL);
    fseek(test_images_fp,0L,SEEK_SET);
    fread(test_images_mem, flen, 1, train_images_fp);
    fclose(test_images_fp);

    FILE *test_labels_fp = NULL;
    test_labels_fp = fopen(test_labels_path, "r");
    fseek(test_labels_fp,0L,SEEK_END);
    flen=ftell(test_labels_fp);
    uint8_t *test_labels_mem = (uint8_t *)malloc(flen+1);
    assert(test_labels_mem != NULL);
    fseek(test_labels_fp,0L,SEEK_SET);
    fread(test_labels_mem, flen, 1, train_images_fp);
    fclose(test_labels_fp);

    srand((unsigned)time(0));

    //uint8_t read_buf[IMAGE_ROW*IMAGE_COL];
    uint64_t image_index[BATCH_SIZE];




    double actual_w[NUM_OUTPUT][NUM_INPUT];

    double delta_w[NUM_OUTPUT][NUM_INPUT];
    double actual_b[NUM_OUTPUT];
    double delta_b[NUM_OUTPUT];
    //double actual_output[NUM_OUTPUT][BATCH_SIZE];
    //double actual_output_softmax[NUM_OUTPUT][BATCH_SIZE];
    double lose[NUM_OUTPUT][BATCH_SIZE];
    for(uint64_t i =0;i<NUM_OUTPUT;i++)
    {
        for(uint64_t j = 0;j<NUM_INPUT;j++)
        {
            actual_w[i][j] = ((double)rand()) / RAND_MAX;
        }
        actual_b[i] = ((double)rand()) / RAND_MAX;
    }
    //print_actual_w((double *)actual_w);

    for(uint64_t epoch = 0;epoch < NUM_EPOCHS;epoch++)
    {
        if(epoch%1000 == 0)
        {
            printf("epoch:%ld\n", epoch);
        }
        for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
        {
            image_index[batch] = rand() % TRAIN_NUM;
            //uint64_t label = get_one_label(train_labels_mem,image_index[batch]);
            //printf("epoch:%ld, batch:%ld, label:%ld, image_index[batch]:%ld, %s\n", epoch, batch, label, image_index[batch], class_name[label]);
        }

        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                lose[output_index][batch] = 0;
                uint8_t *image = get_one_image(train_images_mem, image_index[batch]);
                for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
                {
                    lose[output_index][batch] += actual_w[output_index][input_index] * (image[input_index]);
                }
                lose[output_index][batch] += actual_b[output_index];
                uint64_t label = get_one_label(train_labels_mem,image_index[batch]);
                lose[output_index][batch] -= ((label==output_index)?(500000):(-500000));
                //if(epoch == (NUM_EPOCHS -1))
                //{
                //    printf("lose[%ld][%ld]:%lf\n", output_index, batch, lose[output_index][batch]);
                //}
            }
        }

        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
            {
                delta_w[output_index][input_index] = 0;
                for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
                {
                    uint8_t *image = get_one_image(train_images_mem, image_index[batch]);
                    delta_w[output_index][input_index] += (image[input_index]) * lose[output_index][batch];
                }
                delta_w[output_index][input_index] = delta_w[output_index][input_index] * ETA / BATCH_SIZE / NUM_INPUT / 128;
                //printf("delta_w[%ld][%ld]:%lf\n", output_index, input_index, delta_w[output_index][input_index]);
                actual_w[output_index][input_index] -= delta_w[output_index][input_index];
            }
            delta_b[output_index] = 0;
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                //uint8_t *image = get_one_image(train_images_mem, image_index[batch]);
                delta_b[output_index] += lose[output_index][batch];
            }
            delta_b[output_index] = delta_b[output_index] * ETA / BATCH_SIZE;
            actual_b[output_index] -= delta_b[output_index];

        }
    }

    uint64_t same_num = 0;
    for(uint64_t test_epoch = 0;test_epoch < NUM_EPOCHS_TEST;test_epoch++)
    {
        uint64_t image_test_index;
        image_test_index = rand() % TEST_NUM;

        
        double lose_test[NUM_OUTPUT];
        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            lose_test[output_index] = 0;
            uint8_t *image = get_one_image(test_images_mem, image_test_index);
            for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
            {
                lose_test[output_index] += actual_w[output_index][input_index] * image[input_index];
            }
            lose_test[output_index] += actual_b[output_index];
            //printf("lose_test[%ld]:%lf\n", output_index, lose_test[output_index]);
        }

        uint64_t rv = 0;
        double max = lose_test[0];
        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            if(lose_test[output_index] > max)
            {
                rv = output_index;
                max = lose_test[output_index];
            }
        }
        uint64_t label_test = get_one_label(test_labels_mem,image_test_index);
        if(rv == label_test)
        {
            same_num++;
        }
        //printf("test_epoch:%ld, actural  value:%ld, image_index[batch]:%ld, %s\n",   test_epoch, label_test, image_test_index, class_name[label_test]);
        //printf("test_epoch:%ld, forecast value:%ld, image_index[batch]:%ld, %s\n\n", test_epoch, rv, image_test_index, class_name[rv]);
        //printf("rv[%ld]:%ld\r\n", batch, rv[batch]);
    }
    
    printf("same_num:%ld\n", same_num);
#if 0
    for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
    {
        DIVIDING_LINE;
        printf("output_index:%ld, name:%s\n", output_index, class_name[output_index]);
        for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
        {
            printf("%03.0lf ", 10*actual_w[output_index][input_index]);
            if((input_index + 1) % 28 == 0)
            {
                printf("\n");
            }
        }
    }
#endif
    for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
    {
        printf("%lf\n", actual_b[output_index]);
    }
#if 0
    for(uint64_t epoch = 0;epoch < NUM_EPOCHS;epoch++)
    {
        for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
        {
            image_index[batch] = rand() % TRAIN_NUM;
            uint64_t label = get_one_label(train_labels_fp,image_index[batch]);
            printf("epoch:%ld, batch:%ld, image_index[batch]:%ld, %s\n", epoch, batch, image_index[batch], class_name[label]);
        }

        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            for(uint64_t input_index =0;input_index<2;input_index++)
            {
                delta_w[output_index][input_index] = 0;
                double gradient = 0;
                for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
                {
                    uint8_t image[IMAGE_ROW*IMAGE_COL];
                    get_one_image(train_images_fp, image, sizeof(image), image_index[batch]);
                    gradient = 0;
                    uint64_t input_index_acc =0;
                    for(input_index_acc =0;input_index_acc<NUM_INPUT;input_index_acc++)
                    {
                        gradient = gradient + actual_w[output_index][input_index_acc];//*image[input_index_acc];
                    }
                        printf("input_index_acc:%ld, gradient:%lf\n", input_index_acc, gradient);
#if 0
                    gradient += actual_b[output_index];
                    uint64_t lable = get_one_label(train_labels_fp,image_index[batch]);
                    gradient -= (lable==output_index)?784:0;
                    gradient = gradient * image[input_index];
#endif
                    //printf("gradient:%lf\n", gradient);
                    //delta_w[output_index][input_index] += gradient;
                }
                //delta_w[output_index][input_index] = delta_w[output_index][input_index] * ETA / BATCH_SIZE;
                //actual_w[output_index][input_index] -= delta_w[output_index][input_index];
                printf("\n");
            }
        }

        //print_actual_w((double *)actual_w);

        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            delta_b[output_index] = 0;
            double gradient = 0;
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                uint8_t image[IMAGE_ROW*IMAGE_COL];
                get_one_image(train_images_fp, image, sizeof(image), image_index[batch]);
                for(uint64_t input_index_acc =0;input_index_acc<NUM_INPUT;input_index_acc++)
                {
                    gradient += actual_w[output_index][input_index_acc]*image[input_index_acc];
                }
                gradient += actual_b[output_index];
                uint64_t lable = get_one_label(train_labels_fp,image_index[batch]);
                gradient -= lable==output_index?1:0;
                delta_b[output_index] += gradient;
            }
            delta_b[output_index] = delta_b[output_index] * ETA / BATCH_SIZE;
            actual_b[output_index] -= delta_b[output_index];
        }

        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                actual_output[output_index][batch] = 0;
                uint8_t image[IMAGE_ROW*IMAGE_COL];
                get_one_image(train_images_fp, image, sizeof(image), image_index[batch]);
                for(uint64_t input_index =0;input_index<NUM_INPUT;input_index++)
                {
                    actual_output[output_index][batch] += (image[input_index]*actual_w[output_index][input_index]);
                }
                actual_output[output_index][batch] += actual_b[output_index];
            }
        }
#if 0
        softmax((double *)actual_output, BATCH_SIZE, NUM_OUTPUT, (double *)actual_output_softmax);
        for(uint64_t output_index = 0;output_index<NUM_OUTPUT;output_index++)
        {
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                printf("%1.08f ", actual_output_softmax[output_index][batch]);
            }
            printf("\n");
            for(uint64_t batch = 0;batch < BATCH_SIZE;batch++)
            {
                printf("%f\n", actual_output[output_index][batch]);
            }
        }
        printf("\n");
#endif
    }
    


    #endif
    free(train_images_mem);
    free(train_labels_mem);
    free(test_images_mem);
    free(test_labels_mem);

    FUNC_STOP;
    return 0;
}

int main(int argc, char **argv)
{
    UNUSED(argc);
    UNUSED(argv);
    softmax_new();
    return 0;
}
